# batch_inference_4class.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
VEHICLE_MODEL_PATH = r"D:\CodeCNN\intel\vehicle-detection-0202\FP32\vehicle-detection-0202.xml"
ROAD_MODEL_PATH = r"D:\CodeCNN\intel\road-segmentation-adas-0001\FP32\road-segmentation-adas-0001.xml"
INPUT_FOLDER = r"H:\xiaomi\test"  # 输入图片文件夹
OUTPUT_FOLDER = r"detect"  # 输出结果文件夹

# 创建输出目录
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# -----------------------------
# 初始化 OpenVINO
# -----------------------------
core = Core()
device = "CPU"  # 可选 "GPU"（若支持）

# -----------------------------
# 1. 加载车辆检测模型
# -----------------------------
print("Loading vehicle detection model...")
vehicle_model = core.read_model(model=VEHICLE_MODEL_PATH)
vehicle_compiled = core.compile_model(model=vehicle_model, device_name=device)
vehicle_input_layer = vehicle_compiled.input(0)
vehicle_output_layer = vehicle_compiled.output(0)

# 获取输入尺寸
_, _, h_det, w_det = vehicle_input_layer.shape
print(f"Vehicle detection input shape: {h_det}x{w_det}")

# -----------------------------
# 2. 加载道路分割模型（4类）
# -----------------------------
print("Loading 4-class road segmentation model...")
road_model = core.read_model(model=ROAD_MODEL_PATH)
road_compiled = core.compile_model(model=road_model, device_name=device)
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

# 输出形状应为 [1, 4, H, W]
_, num_classes, h_seg, w_seg = road_output_layer.shape
print(f"Road segmentation output classes: {num_classes} (should be 4)")
if num_classes != 4:
    print("⚠️ 模型输出不是4类！请确认是否使用了多分类版本（如 -1024x768）")
    exit(1)

print(f"Road segmentation input shape: {h_seg}x{w_seg}")

# -----------------------------
# 3. 获取 A 文件夹中所有图片
# -----------------------------
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
image_files = []
for root, dirs, files in os.walk(INPUT_FOLDER):
    for file in files:
        if os.path.splitext(file)[1].lower() in image_extensions:
            image_files.append(os.path.join(root, file))

if not image_files:
    print(f"❌ 没有在 '{INPUT_FOLDER}' 中找到图片！")
    exit(1)

print(f"✅ 发现 {len(image_files)} 张图片，开始批量处理...")

# -----------------------------
# 4. 定义 4 类颜色映射（BGR 格式）- 背景透明
# -----------------------------
# 0: background → 透明（不绘制）
# 1: road → 绿色
# 2: curb → 蓝色（路缘）
# 3: road markings → 红色（车道线）
COLOR_MAP = {
    0: None,  # 背景 - 透明
    1: (0, 255, 0),  # green - 道路
    2: (255, 0, 0),  # blue - 路缘
    3: (0, 0, 255),  # red - 道路标记
}

# -----------------------------
# 5. 批量处理每张图片
# -----------------------------
for i, img_path in enumerate(image_files, 1):
    print(f"\n--- 处理第 {i}/{len(image_files)} 张: {os.path.basename(img_path)} ---")

    # 读取图像
    image = cv2.imread(img_path)
    if image is None:
        print(f"⚠️ 无法加载图像: {img_path}")
        continue

    orig_h, orig_w = image.shape[:2]
    print(f"Original size: {orig_w}x{orig_h}")

    # -----------------------------
    # 6. 车辆检测推理
    # -----------------------------
    resized_det = cv2.resize(image, (w_det, h_det))
    input_det = np.expand_dims(resized_det.transpose(2, 0, 1), axis=0)  # HWC -> NCHW
    detections = vehicle_compiled([input_det])[vehicle_output_layer]

    # -----------------------------
    # 7. 道路分割推理（4类）
    # -----------------------------
    resized_seg = cv2.resize(image, (w_seg, h_seg))
    input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)

    # 推理
    infer_result = road_compiled([input_seg])
    tensor_output = infer_result[road_output_layer]

    # 转为 NumPy 并处理
    seg_probs = np.array(tensor_output)  # shape: [1, 4, H, W]
    seg_class = np.argmax(seg_probs[0], axis=0).astype(np.uint8)
    seg_class = np.ascontiguousarray(seg_class)  # 确保内存连续

    # 调试信息
    print(f"seg_class dtype: {seg_class.dtype}, shape: {seg_class.shape}")
    unique_vals, counts = np.unique(seg_class, return_counts=True)
    print(f"分割结果类别分布:")
    for val, count in zip(unique_vals, counts):
        print(f"  类别 {val}: {count} 像素 ({count / seg_class.size * 100:.2f}%)")

    # -----------------------------
    # 8. 将分割结果缩放到原图尺寸
    # -----------------------------
    seg_class_resized = cv2.resize(seg_class, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)

    # -----------------------------
    # 9. 生成彩色语义分割掩码
    # -----------------------------
    seg_color_mask = np.zeros_like(image)
    non_bg_mask = np.zeros((orig_h, orig_w), dtype=bool)  # 非背景区域掩码

    for cls, color in COLOR_MAP.items():
        if color is None:  # 跳过背景
            continue
        mask = (seg_class_resized == cls)
        mask_pixels = np.sum(mask)
        print(f"类别 {cls} 的像素数: {mask_pixels}")

        if mask_pixels > 0:
            seg_color_mask[mask] = color
            non_bg_mask = non_bg_mask | mask  # 累积非背景区域
            print(f"  → 为类别 {cls} 应用颜色 {color}")

    # 检查背景处理
    bg_mask = (seg_class_resized == 0)
    bg_pixels = np.sum(bg_mask)
    print(f"背景像素数: {bg_pixels}")
    print(f"非背景区域像素数: {np.sum(non_bg_mask)}")

    # -----------------------------
    # 10. 绘制车辆检测框
    # -----------------------------
    output_image = image.copy()
    car_count = 0
    for det in detections[0, 0]:
        conf = float(det[2])
        if conf > 0.5:
            x_min = int(det[3] * orig_w)
            y_min = int(det[4] * orig_h)
            x_max = int(det[5] * orig_w)
            y_max = int(det[6] * orig_h)
            cv2.rectangle(output_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
            cv2.putText(output_image, f"Car: {conf:.2f}", (x_min, y_min - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            car_count += 1

    print(f"检测到车辆数: {car_count}")

    # -----------------------------
    # 11. 叠加语义分割（只在非背景区域叠加）
    # -----------------------------
    overlay = output_image.copy()

    # 只在非背景区域进行叠加
    if np.any(non_bg_mask):
        overlay[non_bg_mask] = cv2.addWeighted(
            output_image[non_bg_mask], 0.6,
            seg_color_mask[non_bg_mask], 0.4, 0
        )
        print("已应用语义分割叠加")
    else:
        print("没有检测到道路相关区域，保持原图")

    # -----------------------------
    # 12. 保存结果到 detect 目录
    # -----------------------------
    base_name = os.path.splitext(os.path.basename(img_path))[0]
    output_filename = f"{base_name}_result.jpg"
    output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    # 如果文件已存在，添加时间戳避免覆盖
    if os.path.exists(output_path):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        output_filename = f"{base_name}_result_{timestamp}.jpg"
        output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    cv2.imwrite(output_path, overlay)
    print(f"✅ 结果已保存至: {output_path}")

    # -----------------------------
    # 13. （可选）显示最后一张图
    # -----------------------------
    if i == len(image_files):  # 只显示最后一张
        def resize_to_fit(img, target_w=1360, target_h=768):
            h, w = img.shape[:2]
            scale = min(target_w / w, target_h / h)
            new_w, new_h = int(w * scale), int(h * scale)
            resized = cv2.resize(img, (new_w, new_h))
            canvas = np.zeros((target_h, target_w, 3), dtype=np.uint8)
            y_pad = (target_h - new_h) // 2
            x_pad = (target_w - new_w) // 2
            canvas[y_pad:y_pad + new_h, x_pad:x_pad + new_w] = resized
            return canvas


        display_img = resize_to_fit(overlay)
        cv2.namedWindow("Batch Result", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Batch Result", 1360, 768)
        cv2.imshow("Batch Result", display_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

print(f"\n🎉 所有 {len(image_files)} 张图片处理完成！结果保存在 '{OUTPUT_FOLDER}' 目录中。")